# Copyright 2023 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Data preprocessing script"""
import os
import json
import logging
import numpy as np
# import utils
import sys
sys.path.append('..')
import utils


class Dataset(object):
    """
    This module implements the APIs for loading dataset and providing batch data
    """
    def __init__(self, args):
        self.logger = logging.getLogger("GraphCM")
        self.max_d_num = args.max_d_num
        self.gpu_num = args.gpu_num
        self.dataset = args.dataset
        self.data_dir = os.path.join('./data', self.dataset)
        self.args = args

        self.train_set = self.load_dataset(os.path.join(self.data_dir, 'train_per_query_quid.txt'), mode='train')
        self.valid_set = self.load_dataset(os.path.join(self.data_dir, 'valid_per_query_quid.txt'), mode='valid')
        self.test_set = self.load_dataset(os.path.join(self.data_dir, 'test_per_query_quid.txt'), mode='test')

        self.trainset_size = len(self.train_set)
        print(self.trainset_size)
        self.validset_size = len(self.valid_set)
        print(self.validset_size)
        self.testset_size = len(self.test_set)

        self.query_qid = utils.load_dict(self.data_dir, 'query_qid.dict')  # 查询或浏览数据
        # 对于每个查询，搜索引擎返回一个排序列表Di,Di有三个属性：唯一的URL标识符，排名位置以及垂直类型
        self.url_uid = utils.load_dict(self.data_dir, 'url_uid.dict')  # 唯一的URL标识符
        self.vtype_vid = utils.load_dict(self.data_dir, 'vtype_vid.dict')  # 垂直类型
        self.query_size = len(self.query_qid)
        self.doc_size = len(self.url_uid)
        self.vtype_size = len(self.vtype_vid)

        self.logger.info('Train set size: {} sessions.'.format(len(self.train_set)))
        self.logger.info('Dev set size: {} sessions.'.format(len(self.valid_set)))
        self.logger.info('Test set size: {} sessions.'.format(len(self.test_set)))
        self.logger.info('Unique query num, including zero vector: {}'.format(self.query_size))
        self.logger.info('Unique doc num, including zero vector: {}'.format(self.doc_size))
        self.logger.info('Unique vtype num, including zero vector: {}'.format(self.vtype_size))

    def load_dataset(self, data_path, mode):
        """
        Loads the dataset
        """
        data_set = []
        lines = open(data_path).readlines()  # 按行读取
        previous_sid = -1
        qids, uids, vids, clicks, relevances = [], [], [], [], []
        for line in lines:
            attr = line.strip().split('\t')  # 删除每一行的空格符（split('\t') ）以及换行符
            sid = int(attr[0].strip())
            if previous_sid != sid:  # 如果不是最后一行
                # a new session starts
                if previous_sid != -1:
                    assert len(uids) // 10 == len(qids)
                    assert len(vids) // 10 == len(qids)
                    assert len(relevances) // 10 == len(qids)
                    assert (len(clicks) - 1) // 10 == len(qids)
                    last_rank = 0
                    for idx, click in enumerate(clicks[1:]):
                        last_rank = idx + 1 if click else last_rank
                    relevance_start = 0
                    for idx, relevance in enumerate(relevances):
                        if relevance != -1:
                            relevance_start = idx
                            assert relevance_start % 10 == 0
                            break
                    data_set.append({'sid': previous_sid,
                                     'qids': qids,
                                     'uids': uids,
                                     'vids': vids,
                                     'clicks': clicks,
                                     'last_rank': last_rank,
                                     'relevances': relevances[relevance_start: relevance_start + 10],
                                     'relevance_start': relevance_start})
                previous_sid = sid
                qids = [int(attr[1].strip())]
                uids = json.loads(attr[2].strip())
                vids = json.loads(attr[3].strip())
                clicks = [0] + json.loads(attr[4].strip())
                relevances = json.loads(attr[5].strip()) if mode == 'label' else [0 for _ in range(self.max_d_num)]
            else:
                # the previous session continues
                qids.append(int(attr[1].strip()))
                uids = uids + json.loads(attr[2].strip())
                vids = vids + json.loads(attr[3].strip())
                clicks = clicks + json.loads(attr[4].strip())
                relevances = relevances + (
                    json.loads(attr[5].strip()) if mode == 'label' else [0 for _ in range(self.max_d_num)])
        last_rank = 0
        for idx, click in enumerate(clicks[1:]):
            last_rank = idx + 1 if click else last_rank
        relevance_start = 0
        for idx, relevance in enumerate(relevances):
            if relevance != -1:
                relevance_start = idx
                assert relevance_start % 10 == 0
                break
        data_set.append({'sid': previous_sid,
                         'qids': qids,
                         'uids': uids,
                         'vids': vids,
                         'clicks': clicks,
                         'last_rank': last_rank,
                         'relevances': relevances[relevance_start: relevance_start + 10],
                         'relevance_start': relevance_start})
        return data_set

    def _one_mini_batch(self, data, indices):
        """
        Get one mini batch data
        """
        batch_data = {'raw_data': [data[i] for i in indices],
                      'qids': [],
                      'uids': [],
                      'vids': [],
                      'clicks': [],
                      'last_ranks': [],
                      'relevances': [],
                      'true_clicks': [],
                      'relevance_starts': []}
        for sidx, sample in enumerate(batch_data['raw_data']):
            batch_data['qids'].append(sample['qids'])
            batch_data['uids'].append(sample['uids'])
            batch_data['vids'].append(sample['vids'])
            batch_data['clicks'].append(sample['clicks'])
            batch_data['last_ranks'].append(sample['last_rank'])
            batch_data['relevances'].append(sample['relevances'])
            batch_data['true_clicks'].append(sample['clicks'][1:])
            batch_data['relevance_starts'].append(sample['relevance_start'])

        q_max_len = max(len(l) for l in batch_data['qids'])
        batch_data['qids'] = list(map(lambda l: l + [0] * (q_max_len - len(l)), batch_data['qids']))
        batch_data['qids'] = np.concatenate(np.expand_dims(batch_data['qids'], axis=0), axis=0)

        u_max_len = max(len(l) for l in batch_data['uids'])
        batch_data['uids'] = list(map(lambda l: l + [0] * (u_max_len - len(l)), batch_data['uids']))
        batch_data['uids'] = np.concatenate(np.expand_dims(batch_data['uids'], axis=0), axis=0)

        v_max_len = max(len(l) for l in batch_data['vids'])
        batch_data['vids'] = list(map(lambda l: l + [0] * (v_max_len - len(l)), batch_data['vids']))
        batch_data['vids'] = np.concatenate(np.expand_dims(batch_data['vids'], axis=0), axis=0)

        c_max_len = max(len(l) for l in batch_data['clicks'])
        batch_data['clicks'] = list(map(lambda l: l + [0] * (c_max_len - len(l)), batch_data['clicks']))
        batch_data['clicks'] = np.concatenate(np.expand_dims(batch_data['clicks'], axis=0), axis=0)
        batch_data['clicks'] = batch_data['clicks'][:, :-1]

        return batch_data

    def gen_mini_batches(self, set_name, batch_size, shuffle=True):
        """
        Generate data batches for a specific dataset (train/valid/test/label)
        """
        if set_name == 'train':
            data = self.train_set
        elif set_name == 'valid':
            data = self.valid_set
        elif set_name == 'test':
            data = self.test_set
        else:
            raise NotImplementedError('No data set named as {}'.format(set_name))
        data_size = len(data)

        indices = np.arange(data_size)  # 大小0到data_size-1
        if shuffle:
            np.random.shuffle(indices)  # 重新排序
        indices = indices.tolist()  # 将数组转为列表

        indices += indices[:(self.gpu_num - data_size % self.gpu_num) % self.gpu_num]
        for batch_start in np.arange(0, len(list(indices)), batch_size):
            batch_indices = indices[batch_start: batch_start + batch_size]
            yield self._one_mini_batch(data, batch_indices)